package cn.fufeiii.bt.demo;

import cn.hutool.core.lang.Snowflake;
import org.hibernate.HibernateException;
import org.hibernate.engine.spi.SharedSessionContractImplementor;
import org.redisson.api.RAtomicLong;
import org.redisson.api.RLock;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.Serializable;
import java.util.concurrent.TimeUnit;

/**
 * 雪花算法ID生成器
 *
 * @author FuFei
 */
@Component
public class SnowflakeIdGenerator implements InitializingBean, com.baomidou.mybatisplus.core.incrementer.IdentifierGenerator, org.hibernate.id.IdentifierGenerator {
    @Autowired
    private RedissonClient redissonClient;

    private Snowflake snowflake;

    @Override
    public void afterPropertiesSet() throws InterruptedException {
        String namespace = "IdGenerator";
        String workIdName = namespace + ":workId";
        String workIdLockName = namespace + ":workId:lock";
        RLock lock = redissonClient.getLock(workIdLockName);
        boolean tryLock = lock.tryLock(10, TimeUnit.SECONDS);
        long workId;
        try {
            if (!tryLock) {
                throw new RuntimeException("初始化SnowflakeIdGenerator失败");
            }
            RAtomicLong atomicLong = redissonClient.getAtomicLong(workIdName);
            workId = atomicLong.get();
            atomicLong.set(workId == 31L ? 0 : workId + 1);
        } finally {
            if (tryLock) {
                lock.unlock();
            }
        }
        snowflake = new Snowflake(workId, 0L);
    }

    @Override
    public Number nextId(Object entity) {
        return snowflake.nextId();
    }

    @Override
    public Serializable generate(SharedSessionContractImplementor sharedSessionContractImplementor, Object o) throws HibernateException {
        return snowflake.nextId();
    }

}
